'''
Author: duliang thinktanker@163.com
Date: 2023-12-17 13:05:57
LastEditors: duliang thinktanker@163.com
LastEditTime: 2025-09-29 14:28:39
# FilePath: 
Description: 
'''
from paddleocr import PaddleOCR
import paddle
import time
import uvicorn
from fastapi import FastAPI, Form, Body, File, Request, Response, HTTPException, status
from fastapi.responses import HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.responses import FileResponse, JSONResponse
from fastapi.templating import Jinja2Templates
import base64
from concurrent.futures import ThreadPoolExecutor
import asyncio
from typing import List
import threading
import queue

MAX_G = 1.5
MAX_MEMORRY_SIZE = 1024 * 1024 * 1024 * MAX_G  # 最大内存限制为3G
# 需要屏蔽的IP列表
# ALLOWED_IPS = ["127.0.0.1", "192.168.15.101", "223.68.11.42"]
print("载入文本识别")
paddle.set_device("gpu:0")
# --cls_model_dir="./inference/ch_ppocr_mobile_v2.0_cls_infer/"     # 分类模型所在文件夹
cls_model_dir = "./whl/cls/ch/ch_ppocr_mobile_v2.0_cls_infer/"
# --det_model_dir="./inference/ch_ppocr_mobile_v2.0_det_infer/"     # 检测模型所在文件夹
det_model_dir = r".\whl\det\ch\ch_PP-OCRv4_det_infer"
# --rec_model_dir="./inference/ch_ppocr_mobile_v2.0_rec_infer/"     # 识别模型所在文件夹
rec_model_dir = r".\whl\rec\ch\ch_PP-OCRv4_rec_infer"

# --use_angle_cls=True --use_space_char=True --use_gpu=False
use_angle_cls = True  # 是否加载分类模型
use_space_char = True  # 是否识别空格
use_gpu = True  # 是否使用GPU，若想利用gpu，设置为True
cls = False  #   前向时是否启动分类
# 实例化PaddleOCR识别类
my_ocr = PaddleOCR(
    cls_model_dir=cls_model_dir,  #
    det_model_dir=det_model_dir,  #
    rec_model_dir=rec_model_dir,  #
    use_angle_cls=use_angle_cls,  #
    # det_max_side_len=680,  #
    # # det_limit_side_len=680,  #
    # max_text_length=25,  #
    # image_shape='320,320',
    # # rec_image_shape='320,320',
    # rec_batch_num=10,
    # lang='ch',
    # # use_tensorrt=True,
    # use_fp16=True,
    # ocr_version='PP-OCRv4',
    use_gpu=use_gpu)

# 创建线程池，预建立4个线程
thread_pool = ThreadPoolExecutor(max_workers=4)
# 创建任务队列
task_queue = queue.Queue()
# 创建线程安全锁
queue_lock = threading.Lock()

# 设置Jinja2模板目录
templates = Jinja2Templates(directory="templates")


def ocr_box(data):
    with open('result.jpg', 'wb') as f:
        f.write(data)
    results = my_ocr.ocr(data, cls=False)
    print(f"识别后显存（阈值：{MAX_G}G）：\033[33m",
          paddle.device.cuda.memory_reserved() // 1024 // 1024, 'M\033[0m')
    if paddle.device.cuda.memory_reserved() > MAX_MEMORRY_SIZE:
        paddle.device.cuda.synchronize(0)
        paddle.device.cuda.empty_cache()
        print(f"清理后：\033[33m",
              paddle.device.cuda.memory_reserved() // 1024 // 1024, 'M\033[0m')
    if results:
        sendContent = []
        sendBox = []
        for result in results:
            if result:
                for t0 in result:
                    # sendContent += ' ' + str(t0[-1][0])
                    sendBox.append(t0[0])
                    sendContent.append(str(t0[-1][0]))
            # print(sendContent)
            return sendContent, sendBox, results
        else:
            return [], [], []
    else:
        return [], [], []


app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    # 允许跨域的源列表，例如 ["http://www.example.org"] 等等，["*"] 表示允许任何源
    allow_origins=["*"],
    # 跨域请求是否支持 cookie，默认是 False，如果为 True，allow_origins 必须为具体的源，不可以是 ["*"]
    allow_credentials=False,
    # 允许跨域请求的 HTTP 方法列表，默认是 ["GET"]
    allow_methods=["*"],
    # 允许跨域请求的 HTTP 请求头列表，默认是 []，可以使用 ["*"] 表示允许所有的请求头
    # 当然 Accept、Accept-Language、Content-Language 以及 Content-Type 总之被允许的
    allow_headers=["*"],
    # 可以被浏览器访问的响应头, 默认是 []，一般很少指定
    # expose_headers=["*"]
    # 设定浏览器缓存 CORS 响应的最长时间，单位是秒。默认为 600，一般也很少指定
    # max_age=1000
)


@app.post("/ocr")
async def read_items(imgbin: bytes = File(...)):
    # 提交任务到线程池进行处理
    loop = asyncio.get_event_loop()
    result = await loop.run_in_executor(thread_pool, ocr_box, imgbin)
    return result[:2]


@app.get("/ocr")
async def get_ocr():
    return 'ok'


def auto_empty_cache():
    print("自动回收显存开启")
    print("每10s施放显存")
    while 1:
        gpu_memory = paddle.device.cuda.memory_reserved() // 1024 // 1024
        if gpu_memory > 3000:
            print('10S清理已用显存（阈值：3G）：', gpu_memory, 'M')
            paddle.device.cuda.empty_cache()
        time.sleep(10)


if __name__ == "__main__":
    uvicorn.run(app="easyocr:app", host='0.0.0.0', port=28211)
